# -*- coding: utf-8 -*-
import requests
import json
import logging
from typing import Any, Dict, List, Text, Optional

logger = logging.getLogger(__name__)


class QwenAPI:
    """Class to connect to the Rasa X API"""

    def __init__(self ):
        self.url = 'http://localhost:11434/api/chat'
        self.device = "cpu" # the device = cuda


    def predict(self, msg):
        logger.info('#### predict='+ msg )
        # headers = {"Content-Type": "application/json; charset=utf-8"}
        # ret = requests.post(url='http://localhost:9090/hi', headers=headers)
        # print(ret.text)

        params = { "model": "qwen:0.5b-chat",   "messages": [     {       "role": "user",       "content": msg     }   ],   "stream": False }
        json_data = json.dumps(params)
        # print(json_data)
        headers = {"Content-Type": "application/json; charset=utf-8"}
        ret = requests.post(url=self.url, data=json_data, headers=headers)
        print(ret.text)
        json_data = json.loads(ret.text)
        return json_data['message']['content']

qwen = QwenAPI()
ret = qwen.predict('你好')
print(ret)